#ifndef __MODEL_RUNNER_H_
#define __MODEL_RUNNER_H_

#include <chrono>
#include <memory>
#include <set>


#include "model.h"

class ModelRunner {
 public:
  ModelRunner() {

  }

  int Init(std::shared_ptr<ModelLoader> model) {
    model_ = model;
    if (!model_->Init()) {
      model_->Init();
    }
    return 0;
  }

  int CreateBuffer(uint64_t *buff_handle) {
    aipu_status_t ret;
    aipu_create_job_cfg create_job_cfg;
    ret = aipu_create_job(model_->GetContext(), model_->GetGraphId(), buff_handle, &create_job_cfg);
    if (ret < 0) {
      *buff_handle = 0;
      std::cout << "create buffer failed" << std::endl;
      return -1;
    }
    handle_set_.insert(*buff_handle);
    return 0;
  }

  int Run(uint64_t buffer_handle, std::vector<char*> inputs) {
    aipu_status_t ret;
    const char* msg = nullptr;
    if (handle_set_.find(buffer_handle) == handle_set_.end()) {
      std::cout << "wrong buffer handle" << std::endl;
      return -1;
    }

    if (inputs.size() != model_->GetInputCnt()) {
      std::cout << "input shape is not same as model input shape, input size: " << inputs.size()
                << ", model_input size: " << model_->GetInputCnt() << std::endl;
      return -1;
    }

    // auto start = std::chrono::high_resolution_clock::now();
    for (size_t i = 0; i < inputs.size(); i++) {
      ret = aipu_load_tensor(model_->GetContext(), buffer_handle, i, inputs[i]);
      if (ret != AIPU_STATUS_SUCCESS) {
        aipu_get_error_message(model_->GetContext(), ret, &msg);
        AIPU_ERR()("aipu_load_tensor: %s\n", msg);
        return -1;
      }
    }
    // auto end = std::chrono::high_resolution_clock::now();
    // std::cout << "total time: " << std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() << "us" << std::endl;
    ret = aipu_finish_job(model_->GetContext(), buffer_handle, -1);
    if (ret != AIPU_STATUS_SUCCESS) {
      aipu_get_error_message(model_->GetContext(), ret, &msg);
      AIPU_ERR()("aipu_finish_job: %s\n", msg);
      return -1;
    }

    return 0;
  }

  ~ModelRunner() {
    aipu_status_t ret;
    const char* msg = nullptr;
    std::cout << "handle set size: " << handle_set_.size() << std::endl;
    for (auto handle : handle_set_) {
      ret = aipu_clean_job(model_->GetContext(), handle);
      if (ret != AIPU_STATUS_SUCCESS) {
        aipu_get_error_message(model_->GetContext(), ret, &msg);
        AIPU_ERR()("aipu_clean_job: %s\n", msg);
      }
    }
  }
 private:
  uint64_t job_id_;
  std::shared_ptr<ModelLoader> model_ = nullptr;
  std::set<uint64_t> handle_set_;
};


#endif
